class ConsistentHash:
    def __init__(self, nodes=None, replicas=100):
        self.replicas = replicas  # 虚拟节点数
        self.ring = {}  # 哈希环
        self.sorted_keys = []  # 排序的节点哈希值
        
        if nodes:
            for node in nodes:
                self.add_node(node)
    
    # def add_node(self, node):
    #     for i in range(self.replicas):
    #         key = self.hash(f"{node}:{i}")
    #         self.ring[key] = node
    #         self.sorted_keys.append(key)
    #     self.sorted_keys.sort()
    
    def add_node(self, node):
        if node in self.ring.values():  # 检查节点是否已存在
            return
        for i in range(self.replicas):
            key = self.hash(f"{node}:{i}")
            self.ring[key] = node
            self.sorted_keys.append(key)
        self.sorted_keys.sort()

    # def remove_node(self, node):
    #     for i in range(self.replicas):
    #         key = self.hash(f"{node}:{i}")
    #         del self.ring[key]
    #         self.sorted_keys.remove(key)

    def remove_node(self, node):
        # 首先检查节点是否存在
        if node not in self.ring.values():
            raise ValueError(f"Node {node} not in ring")
        
        # 收集所有要删除的键
        keys_to_remove = [key for key, n in self.ring.items() if n == node]
        
        if not keys_to_remove:
            raise ValueError(f"No virtual nodes found for {node}")
        
        # 执行删除操作
        for key in keys_to_remove:
            del self.ring[key]
            self.sorted_keys.remove(key)
    
    def get_node(self, key):
        if not self.ring:
            return None
        hash_key = self.hash(key)
        # 顺时针查找第一个节点
        for node_key in self.sorted_keys:
            if hash_key <= node_key:
                return self.ring[node_key]
        return self.ring[self.sorted_keys[0]]
    
    def hash(self, key):
        # 实现一个哈希函数，返回32位整数
        pass



import unittest
from hashlib import md5

class TestConsistentHash(unittest.TestCase):
    def setUp(self):
        # 使用简单的md5哈希作为测试用哈希函数
        self.hash_func = lambda key: int(md5(str(key).encode()).hexdigest(), 16)
        
        # 创建一致性哈希实例，增加虚拟节点数
        self.ch = ConsistentHash(replicas=100)
        self.ch.hash = self.hash_func
        
        # 测试用节点
        self.nodes = ['node1', 'node2', 'node3', 'node4']
        
    def test_empty_ring(self):
        """测试空环的情况"""
        self.assertIsNone(self.ch.get_node('any_key'))
        
    def test_single_node(self):
        """测试单个节点的情况"""
        self.ch.add_node(self.nodes[0])
        self.assertEqual(self.ch.get_node('key1'), self.nodes[0])
        self.assertEqual(self.ch.get_node('key2'), self.nodes[0])
        
    def test_add_nodes(self):
        """测试添加多个节点"""
        for node in self.nodes:
            self.ch.add_node(node)
            
        # 验证节点数量 (节点数 × 虚拟节点数)
        self.assertEqual(len(self.ch.sorted_keys), len(self.nodes) * 100)
        
    def test_key_distribution(self):
        """测试键的分布情况"""
        for node in self.nodes:
            self.ch.add_node(node)
            
        distribution = {node: 0 for node in self.nodes}
        test_keys = [f"key{i}" for i in range(1000)]
        
        for key in test_keys:
            node = self.ch.get_node(key)
            distribution[node] += 1
            
        # 验证所有键都被分配到了节点上
        self.assertEqual(sum(distribution.values()), len(test_keys))
        
        # 放宽均衡检查标准
        avg = len(test_keys) / len(self.nodes)
        for count in distribution.values():
            self.assertTrue(avg * 0.5 < count < avg * 1.5, 
                          f"Distribution not balanced: {distribution}")
        print(f"Actual distribution: {distribution}")
    
    def test_node_removal(self):
        """测试节点移除"""
        for node in self.nodes:
            self.ch.add_node(node)
            
        # 记录原始映射关系
        original_mapping = {}
        test_keys = [f"key{i}" for i in range(100)]
        for key in test_keys:
            original_mapping[key] = self.ch.get_node(key)
            
        # 移除一个节点
        removed_node = self.nodes[1]
        self.ch.remove_node(removed_node)
        
        # 验证剩余的节点
        remaining_nodes = set(self.nodes) - {removed_node}
        self.assertEqual(set(self.ch.ring.values()), remaining_nodes)
        
        # 验证键的重新映射
        migrated = 0
        for key in test_keys:
            new_node = self.ch.get_node(key)
            if new_node != original_mapping[key]:
                migrated += 1
                # 验证重新映射的键确实原来属于被移除的节点
                if original_mapping[key] != removed_node:
                    self.assertIn(new_node, remaining_nodes)
        
        # 验证只有受影响的部分键被重新映射
        self.assertTrue(migrated > 0, "No keys migrated after node removal")
        print(f"Keys migrated after node removal: {migrated}/{len(test_keys)}")
    
    def test_add_remove_sequence(self):
        """测试节点的添加和移除序列"""
        self.ch.add_node(self.nodes[0])
        node_for_key = self.ch.get_node('test_key')
        
        # 添加新节点
        self.ch.add_node(self.nodes[1])
        new_node_for_key = self.ch.get_node('test_key')
        
        # 键可能保持不变或迁移到新节点
        self.assertTrue(new_node_for_key in [self.nodes[0], self.nodes[1]])
        
        # 移除原始节点
        self.ch.remove_node(self.nodes[0])
        final_node_for_key = self.ch.get_node('test_key')
        self.assertEqual(final_node_for_key, self.nodes[1])
    
    def test_virtual_nodes(self):
        """测试虚拟节点的影响"""
        # 创建两个不同虚拟节点数的环
        ch_low = ConsistentHash(replicas=10)
        ch_low.hash = self.hash_func
        ch_high = ConsistentHash(replicas=100)
        ch_high.hash = self.hash_func
        
        for node in self.nodes:
            ch_low.add_node(node)
            ch_high.add_node(node)
            
        # 测试键的分布均匀性
        test_keys = [f"key{i}" for i in range(1000)]
        
        def get_distribution(ch):
            dist = {node: 0 for node in self.nodes}
            for key in test_keys:
                dist[ch.get_node(key)] += 1
            return dist
        
        low_dist = get_distribution(ch_low)
        high_dist = get_distribution(ch_high)
        
        # 计算标准差
        def calc_stddev(dist):
            avg = sum(dist.values()) / len(dist)
            variance = sum((x - avg)**2 for x in dist.values()) / len(dist)
            return variance ** 0.5
        
        stddev_low = calc_stddev(low_dist)
        stddev_high = calc_stddev(high_dist)
        
        # 高虚拟节点数应该有更均匀的分布
        self.assertLess(stddev_high, stddev_low,
                      "Higher replicas should provide better distribution")
    
    def test_consistent_hashing_property(self):
        """测试一致性哈希的核心特性"""
        # 初始节点
        initial_nodes = self.nodes[:2]
        for node in initial_nodes:
            self.ch.add_node(node)
            
        # 记录初始映射
        test_keys = [f"key{i}" for i in range(100)]
        initial_mapping = {key: self.ch.get_node(key) for key in test_keys}
        
        # 添加新节点
        new_node = self.nodes[2]
        self.ch.add_node(new_node)
        
        # 验证最小化重新映射
        migrated = 0
        for key in test_keys:
            if self.ch.get_node(key) != initial_mapping[key]:
                migrated += 1
                # 验证只有部分键被重新映射
                self.assertTrue(migrated < len(test_keys) * 0.5, 
                              "Too many keys remapped after adding one node")
        
        print(f"Keys migrated after adding node: {migrated}/{len(test_keys)}")
    
    def test_edge_cases(self):
        """测试边界情况"""
        # 重置为低虚拟节点数以测试重复添加
        self.ch = ConsistentHash(replicas=3)
        self.ch.hash = self.hash_func
        
        # 测试空键
        self.ch.add_node(self.nodes[0])
        self.assertEqual(self.ch.get_node(''), self.nodes[0])
        
        # 测试非常长的键
        long_key = 'x' * 1000
        self.assertEqual(self.ch.get_node(long_key), self.nodes[0])
        
        # 测试非字符串键
        self.assertEqual(self.ch.get_node(123), self.nodes[0])
        self.assertEqual(self.ch.get_node(('tuple', 'key')), self.nodes[0])
        
        # 测试重复添加节点
        self.ch.add_node(self.nodes[0])
        self.assertEqual(len(self.ch.sorted_keys), 3)  # 应该不增加
        
        # 测试移除不存在的节点
        with self.assertRaises(ValueError):
            self.ch.remove_node('nonexistent_node')

if __name__ == '__main__':
    unittest.main(verbosity=2)



